import tiktoken
import torch  # PyTouch 深度学习框架
from torch.utils.data import Dataset, DataLoader


class GPTDataSetV1(Dataset):
    """自定义 PyTorch Dataset，用于处理 GPT 风格的训练数据"""

    def __init__(self, txt, tokenizer, max_length, stride):
        """
        初始化数据集：
        - txt: 输入的原始文本
        - tokenizer: 分词器（这里使用 tiktoken）
        - max_length: 每个输入序列的最大长度，每个样本的 token 数量
        - stride: 滑动窗口的步长（控制数据重叠程度）（如 1 表示重叠采样，max_length 表示不重叠）。
        """
        self.input_ids = []  # 存储输入序列（token IDs）
        self.target_ids = []  # 存储目标序列（比输入序列右移 1 个 token）

        # 1. 对整个文本进行分词，得到 token IDs（例如 [123, 456, 789, ...]）
        token_ids = tokenizer.encode(txt)
        # 2. 使用滑动窗口生成训练样本（input 和 target）
        for i in range(0, len(token_ids) - max_length, stride):
            # 当前窗口的输入（从 i 到 i+max_length）
            input_chunk = token_ids[i: i + max_length]
            # 当前窗口的目标（从 i+1 到 i+max_length+1，即输入右移 1 位）
            target_chunk = token_ids[i + 1: i + max_length + 1]
            # 将输入和目标转换为 PyTorch Tensor 并存储
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        """根据索引返回单个样本（input 和 target）"""
        return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v1(txt, batch_size=4, max_length=256,
                         stride=128, shuffle=True, drop_last=True,
                         num_workers=0):
    """
    创建 PyTorch DataLoader，用于批量加载数据：
    - txt: 输入的原始文本
    - batch_size: 每批次的样本数
    - max_length: 输入序列的最大长度，每个样本的 token 数量
    - stride: 滑动窗口的步长（如 1 表示重叠采样，max_length 表示不重叠）。
    - shuffle: 是否打乱数据顺序（测试时关闭打乱，训练时建议 True。）
    - drop_last: 是否丢弃最后一个不完整的批次
    - num_workers: 数据加载的线程数（0 表示主线程加载）
    """
    # 1. 初始化分词器（使用 GPT-2 的分词方式）
    tokenizer = tiktoken.get_encoding('gpt2')

    # 2. 创建自定义 Dataset
    dataset = GPTDataSetV1(txt, tokenizer, max_length, stride)

    # 3. 创建 DataLoader，用于批量加载数据， 将 Dataset 转换为 DataLoader，支持批量加载和数据打乱。
    dataloader = DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers)
    return dataloader


# ----------------------------------------------------------------------
# 小型示例（词汇表大小=6，嵌入维度=3）
# 假设我们有一张仅包含6个单词的小型词汇表（而非BPE分词器中包含50257个单词的词汇表），并且想要创建维度为3的嵌入
# 利用vocab_size和output_dim在PyTorch中实例化一个嵌入层。此外，为了确保结果的可重复性，可以将随机种子设置为123
vocab_size = 6
output_dim = 3
torch.manual_seed(123)
# 根据词元size和维度生成权重矩阵
embedding_layer = torch.nn.Embedding(vocab_size, output_dim)
print(embedding_layer.weight)
print(embedding_layer(torch.tensor([3])))

# 词元ID向量
input_ids = torch.tensor([2, 3, 5, 1])
# 从权重矩阵取出对应次元ID的权重向量
print(embedding_layer(input_ids))


# ----------------------------------------------------------------------
# 实际使用（GPT-2 的词汇表大小=50257，嵌入维度=256）
# 词汇量为50257的BPE分词器创建
vocab_size = 50257
output_dim = 256
token_embedding_layer = torch.nn.Embedding(num_embeddings=vocab_size, embedding_dim=output_dim)

# 使用上述的token_embedding_layer，当我们从数据加载器中采样数据时，每个批次中的每个词元都将被嵌入为一个256维的向量。
# 如果设定批次大小为8，且每个批次包含4个词元，则结果将是一个8×4×256的张量。
#
with open('../data/the-verdict.txt', 'r', encoding='utf-8') as f:
    raw_text = f.read()

max_length = 4
downloader = create_dataloader_v1(
    raw_text, batch_size=8, max_length=max_length,
    stride=max_length, shuffle=False
)
data_iter = iter(downloader)
inputs, targets = next(data_iter)
print("Token IDs:\n", inputs)
print("\nInputs shape:\n", inputs.shape)
# Token IDs:
#  tensor([[   40,   367,  2885,  1464],
#         [ 1807,  3619,   402,   271],
#         [10899,  2138,   257,  7026],
#         [15632,   438,  2016,   257],
#         [  922,  5891,  1576,   438],
#         [  568,   340,   373,   645],
#         [ 1049,  5975,   284,   502],
#         [  284,  3285,   326,    11]])
#
# Inputs shape:
#  torch.Size([8, 4])
# 词元ID张量的维度为8×4，这表明数据批次包含8个文本样本，每个样本由4个词元组成。


# 使用嵌入层将这些词元ID嵌入256维的向量中
token_embeddings = token_embedding_layer(inputs)
print(token_embeddings.shape)
# torch.Size([8, 4, 256])

# 生成位置嵌入向量
context_len = max_length
pos_embedding_layer = torch.nn.Embedding(context_len, output_dim)
pos_embeddings = pos_embedding_layer(torch.arange(context_len))
print(pos_embeddings.shape)
# torch.Size([4, 256])

# 位置嵌入向量加入到词元嵌入向量变为大模型需要的输入嵌入向量
input_embeddings = token_embeddings + pos_embeddings
# 最终输入 = 词嵌入 + 位置嵌入。
print(input_embeddings.shape)
# torch.Size([8, 4, 256])


"""
由于大语言模型无法直接处理原始文本，因此我们必须将文本数据转换为名为“嵌入”的数值向量。嵌入将离散的数据（如词语或图像）映射到连续的向量空间，使其能够用于神经网络的训练。
1.首先，原始文本被分解为词元，这些词元可能是单词或字符。然后，这些词元被转换为整数表示，即词元ID。
2.为了增强模型对不同上下文的理解和处理能力，可以引入诸如<|unk|>、<|endoftext|>等特殊词元，来处理未知词汇或在不相关的文本之间划分边界。
3.通过使用BPE分词器，GPT-2、GPT-3等大语言模型可以将未知词汇分解为子词单元或单个字符来有效地处理它们。
4.通过使用滑动窗口方法对已经分词的数据进行采样，可以生成大语言模型训练所需的输入-目标对。
5.作为一种查找操作，PyTorch中的嵌入层用来检索与词元ID对应的向量。所得的嵌入向量为词元提供了连续的表示形式，这对于训练像大语言模型这样的深度学习模型至关重要。
6.尽管词元嵌入为每个词元提供了一致的向量表示，但它们并不包含词元在序列中的位置信息。为了解决这一问题，通常采用两种位置嵌入策略：绝对位置嵌入和相对位置嵌入。
    OpenAI的GPT模型采用了绝对位置嵌入，这些嵌入被添加到词元嵌入向量中，并在模型训练过程中进行优化。
"""


"""
5.整体流程总结
文本 → Token IDs
    使用 tiktoken 分词（如 "hello" → [31373]）。
滑动窗口生成样本
    输入：[token_1, token_2, token_3]
    目标：[token_2, token_3, token_4]。
转换为 DataLoader
    批量加载（如 batch_size=8）。
Token → 嵌入向量
    通过 nn.Embedding 映射到 256 维。
添加位置信息
    位置嵌入 + 词嵌入 = 最终输入。
"""

"""
6. 关键概念
    BPE 分词：将未知词拆分为子词（如 "unhappy" → "un" + "happy"）。
    嵌入层：类似字典查找，将离散 ID 映射为连续向量。
    位置嵌入：让模型感知 token 的顺序（绝对 or 相对位置）。
"""

"""
下一步学习建议
    尝试修改 max_length 和 stride，观察数据如何变化。
    了解 Transformer 模型 如何处理这些嵌入输入。
    扩展知识：相对位置编码（如 RoPE） vs 绝对位置编码。
"""
